#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# # author : cypro666
# # date   : 2015.08.01
# # wrapper of NMF method in sklearn
import sys, math, json
import numpy as np
from threading import Thread
from sklearn import decomposition

from magic3.utils import Timer
from magic3.filesystem import *

add_sys_path(grand_dir(__file__))

from skt.base import MethodBase
from skt.utils import _normalize, make_options

class NMF(MethodBase):
    def __init__(self, parameters={}):
        super().__init__(parameters)

    @property
    def name(self): return 'NMF'

    def output(self, model, weights):
        self.Log('output')
        super().output()
        
        extra = {
            "weights" : [list(a) for a in weights],
            "reconstruction_err" : float(model.reconstruction_err_),
            "n_iter"  : float(model.n_iter_)
        }
        
        fn = self.json_name()
        json.dump(extra, open(fn, 'w'), indent=4)

    def execute(self):
        self.read_input()
        self._normalize(False, self._param['normalize'])

        if self._param['tol'] < 5E-6:
            self._param['tol'] = 5E-6
        
        if self._param['n_components'] < 2:
            self._param['n_components'] = 2
        
        if self._param['nls_max_iter'] < 10:
            self._param['nls_max_iter'] = 10
        
        if self._param['max_iter'] < 10:
            self._param['max_iter'] = 10

        self.save_parameters()

        model = decomposition.NMF(beta=None,
                                  eta=None,
                                  init=self._param['init'],
                                  tol=self._param['tol'],
                                  n_components=self._param['n_components'],
                                  max_iter=self._param['max_iter'],
                                  nls_max_iter=self._param['nls_max_iter'],
                                  random_state=False,
                                  sparseness=None)

        weights = model.fit_transform(self._train)

        self._results = model.components_
        self.output(model, weights)

    def run(self, timeout):
        t = Thread(target=self.execute)
        t.start()
        t.join(timeout)
        if t.is_alive():
            self.Log('timeout!')
        self.Log('exit')


if __name__ == '__main__':
    opts = [('train_file', 'str', []),
            ('results_file', 'str', []),
            ('log_file', 'str', []),
            ('normalize', 'choice', ['0', '1', '2']),
            ('init', 'choice', ['nndsvd', 'nndsvda', 'nndsvdar', 'random']),
            ('n_components', 'int', []),
            ('max_iter', 'int', []),
            ('nls_max_iter', 'int', []),
            ('tol', 'float', [])]

    parameters = make_options(opts)

    timer = Timer()
    NMF(parameters).run(3600)

    if __debug__:
        print('elapsed:', timer)


